#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2021/11/8 15:32
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : customized_k_means.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import ConvexHull
import pandas as pd


class KMeans:
    def __init__(self, k):
        self.k = k

    def fit(self, x, initial_centroid_index=None, max_iters=10, seed=16, plt_process=False):
        global idx
        m, n = x.shape
        # 没有指定中心点时，随机初始化中心点
        if initial_centroid_index is None:
            np.random.seed(seed)
            initial_centroid_index = np.random.randint(0, m, self.k)
        centroid = x[initial_centroid_index, :]

        # 打开交互模式
        plt.ion()
        for i in range(max_iters):
            # 按照中心点给样本分类
            idx = self.find_closest_centroids(x, centroid)
            if plt_process:
                self.plot_converge(x, idx, initial_centroid_index)
            # 重新计算中心点
            centroid = self.compute_centroids(x, idx)
        else:
            # 关闭交互模式
            plt.ioff()
            plt.show()
            return centroid, idx

    @staticmethod
    def find_closest_centroids(x, centroid):
        # 这种方式利用 numpy 的广播机制，直接计算样本到各中心的距离，不用循环，速度比较快，但是在样本比较大时，更消耗内存
        distance = np.sum((x[:, np.newaxis, :] - centroid) ** 2, axis=2)
        idx = distance.argmin(axis=1)
        return idx

    def compute_centroids(self, x, idx):
        centroids = np.zeros((self.k, x.shape[1]))
        for i in range(self.k):
            centroids[i, :] = np.mean(x[idx == i], axis=0)
        return centroids

    def plot_converge(self, x, idx, initial_idx):
        plt.cla()  # 清除原有图像
        plt.title("k-meas converge process")
        plt.xlabel('density')
        plt.ylabel('sugar content')
        plt.scatter(x[:, 0], x[:, 1], c='lightcoral')
        # 标记初始化中心点
        plt.scatter(x[initial_idx, 0], x[initial_idx, 1], label='initial center', c='k')
        # 画出每个簇的凸包
        for i in range(self.k):
            X_i = x[idx == i]
            # 获取当前簇的凸包索引
            hull = ConvexHull(X_i).vertices.tolist()
            hull.append(hull[0])
            plt.plot(X_i[hull, 0], X_i[hull, 1], 'c--')
        plt.legend()
        plt.pause(0.5)


def load_data(data_path, delimiter=None, attr_start_col=1, reindex=None, test_size=0.2):
    """
    加载数据集并且划分特征集和标记集
    :param data_path: 数据集文件路径
    :param delimiter: 分隔符
    :param attr_start_col:开始索引
    :param reindex: 重定义列索引，元组
    :param test_size: 划分测试集比例
    :return:
    """
    from sklearn.model_selection import train_test_split
    dataset = pd.read_csv(data_path, header=0, delimiter=delimiter, names=list(range(*reindex)))
    attr_data = np.array(dataset.iloc[:, attr_start_col:])
    return train_test_split(attr_data, test_size=test_size, random_state=1)


if __name__ == '__main__':
    path = "../dataset/watermelon_4.0.csv"
    attr_train, attr_test = load_data(path, delimiter=",", reindex=(1, 4))
    params = [[2, 12], [3, 6], [4, 23]]
    for k, index in params:
        centroid, idx = KMeans(k).fit(attr_train, initial_centroid_index=index, seed=24)
        print("质心=\n{},\n索引=\n{}".format(centroid, idx))
